// import nj from 'numjs';
import * as tf from "@tensorflow/tfjs";

export async function postprocess_boxes(bbox, score_threshold = 0.3) {
  // const valid_scale = tf.tensor([0, Number.MAX_SAFE_INTEGER], [batch, 2]);
  // const bbox = tf.concat(pred_bbox.map((box) => {
  //   return tf.reshape(box, [-1, box.shape[box.shape.length - 1]]);
  // }))
  const [d1, d2] = bbox.shape;
  const pred_xywh = bbox.slice([0, 0], [-1, 4]);
  const pred_conf = bbox.slice([0, 4], [-1, 1]).squeeze();
  const pred_prob = bbox.slice([0, 5], [-1, d2 - 5])
  // pred_coor = np.concatenate([pred_xywh[:, :2] - pred_xywh[:, 2:] * 0.5,
  // pred_xywh[:, :2] + pred_xywh[:, 2:] * 0.5], axis=-1)
  let pred_coor = tf.concat([tf.sub(pred_xywh.slice([0, 0], [-1, 2]), pred_xywh.slice([0, 2], [-1, 2]).mul(0.5)),
  tf.add(pred_xywh.slice([0, 0], [-1, 2]), pred_xywh.slice([0, 2], [-1, 2]).mul(0.5))], -1);

  // pred_coor = np.concatenate([np.maximum(pred_coor[:, :2], [0, 0]),
  // np.minimum(pred_coor[:, 2:], [org_w - 1, org_h - 1])], axis=-1)
  pred_coor = tf.concat([tf.maximum(pred_coor.slice([0, 0], [-1, 2]), 0),
  tf.minimum(pred_coor.slice([0, 2], [-1, 2]), 416)], -1);
  // invalid_mask = np.logical_or((pred_coor[:, 0] > pred_coor[:, 2]), (pred_coor[:, 1] > pred_coor[:, 3]))
  const invalid_mask = tf.logicalOr(tf.greater(pred_coor.slice([0, 0], [-1, 1]),
    pred_coor.slice([0, 2], [-1, 1])), tf.greater(pred_coor.slice([0, 1], [-1, 1]), pred_coor.slice([0, 3], [-1, 1])))
    .reshape([-1])
  // pred_coor[invalid_mask] = 0
  pred_coor = pred_coor.where(invalid_mask.logicalNot(), 0);


  // bboxes_scale = np.sqrt(np.multiply.reduce(pred_coor[:, 2:4] - pred_coor[:, 0:2], axis=-1))

  const bboxes_scale = tf.sqrt(tf.sub(pred_coor.slice([0, 2], [-1, 2]), pred_coor.slice([0, 0], [-1, 2])).prod(-1))

  // scale_mask = np.logical_and((valid_scale[0] < bboxes_scale), (bboxes_scale < valid_scale[1]))

  const scale_mask = tf.logicalAnd(tf.less(0, bboxes_scale), tf.less(bboxes_scale, Number.MAX_SAFE_INTEGER))
    .reshape([-1])

  // # 5. discard boxes with low scores
  // classes = np.argmax(pred_prob, axis=-1)

  const _classes = tf.argMax(pred_prob, -1);
  // scores = pred_conf * pred_prob[np.arange(len(pred_coor)), classes]

  const _scores = tf.mul(pred_conf, tf.gatherND(pred_prob, tf.concat([tf.range(0, _classes.shape[0], 1, 'int32').expandDims(1), _classes.expandDims(1)], -1)));
  // score_mask = scores > score_threshold

  const score_mask = tf.greater(_scores, score_threshold);
  // mask = np.logical_and(scale_mask, score_mask)
  const mask = tf.logicalAnd(scale_mask, score_mask).reshape([-1]);
  // coors, scores, classes = pred_coor[mask], scores[mask], classes[mask]

  const coors = await tf.booleanMaskAsync(pred_coor, mask);
  const scores = await tf.booleanMaskAsync(_scores, mask);
  const classes = await tf.booleanMaskAsync(_classes, mask);
  // return np.concatenate([coors, scores[:, np.newaxis], classes[:, np.newaxis]], axis=-1)

  return tf.concat([coors, scores.expandDims(1), classes.expandDims(1)], -1)
}

export async function nms(bboxes, iou_threshold = 0.45) {
  const classes_in_img = new Set(bboxes.slice([0, 5], [-1, 1]).squeeze().dataSync())
  const best_bboxes = [];
  let resolve;
  const res = new Promise((_resolve) => {
    resolve = _resolve;
  });

  classes_in_img.forEach(async (cls, i) => {
    const cls_mask = tf.equal(bboxes.slice([0, 5], [-1, 1]).squeeze(), cls);
    let cls_bboxes = await tf.booleanMaskAsync(bboxes, cls_mask);
    while (cls_bboxes.size) {
      const max_ind = tf.argMax(cls_bboxes.slice([0, 4], [-1, 1]).squeeze());
      const best_bbox = tf.gatherND(cls_bboxes, max_ind.expandDims());
      best_bboxes.push(best_bbox);
      const maxInx = max_ind.dataSync()[0];
      cls_bboxes = tf.concat([cls_bboxes.slice([0], [maxInx]), cls_bboxes.slice([maxInx + 1])]);
      const iou = bbxes_iou(best_bbox.slice([0], [4]).expandDims(), cls_bboxes.slice([0, 0], [-1, 4]))
      let weight = tf.ones([iou.size], 'float32');

      const iou_mask = tf.greater(iou, iou_threshold);
      weight = weight.where(iou_mask.logicalNot(), 0.0);

      cls_bboxes = tf.concat([cls_bboxes.slice([0, 0], [-1, 4]), cls_bboxes.slice([0, 4], [-1, 1]).mul(weight), cls_bboxes.slice([0, 5])], -1);
      const score_mask = tf.greater(cls_bboxes.slice([0, 4], [-1, 1]), 0).squeeze(-1);
      cls_bboxes = await tf.booleanMaskAsync(cls_bboxes, score_mask);
    }
    if (i === classes_in_img.size - 1) {
      resolve(best_bboxes)
    }
  })
  // return best_bboxes;
  return res;
}


function bbxes_iou(boxes1, boxes2) {
  const boxes1_area = tf.mul(tf.sub(boxes1.slice([0, 2], [-1, 1]), boxes1.slice([0, 0], [-1, 1])),
    tf.sub(boxes1.slice([0, 3], [-1, 1]), boxes1.slice([0, 1], [-1, 1])))

  const boxes2_area = tf.mul(tf.sub(boxes2.slice([0, 2], [-1, 1]), boxes2.slice([0, 0], [-1, 1])),
    tf.sub(boxes2.slice([0, 3], [-1, 1]), boxes2.slice([0, 1], [-1, 1])))

  const left_up = tf.maximum(boxes1.slice([0, 0], [-1, 2]), boxes2.slice([0, 0], [-1, 2]));
  const right_down = tf.maximum(boxes1.slice([0, 2], [-1, 2]), boxes2.slice([0, 2], [-1, 2]))

  const inter_section = tf.maximum(tf.sub(right_down, left_up), 0);

  const inter_area = tf.mul(inter_section.slice([0, 0], [-1, 1]), inter_section.slice([0, 1], [-1, 1]));

  const union_area = tf.add(boxes1_area, tf.sub(boxes2_area, inter_area));

  return tf.maximum(tf.div(inter_area, union_area), Number.EPSILON).squeeze(-1);


}